package ca.compsci.opent.codegen;

import static org.objectweb.asm.Opcodes.*;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.Channels;
import java.nio.channels.WritableByteChannel;

import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Label;
import org.objectweb.asm.MethodVisitor;

import ca.compsci.opent.analysis.*;
import ca.compsci.opent.node.*;
import ca.compsci.opent.semantics.*;

public class BytecodeGenerator extends DepthFirstAdapter{
	private FileOutputStream out;
	private ClassWriter cw;
	private MethodVisitor mv;
	private SymbolTable table;
	private int max_stack_size = 0, max_locals_size = 0;
	private TypedSymbol.Type last_type = TypedSymbol.Type.DEFFERED;
	private boolean toLong = false;
	
	public BytecodeGenerator(File dir, SymbolTable table) throws FileNotFoundException {
		if ((this.table = table) == null)
			throw new NullPointerException("No symbol table provided");
		if (dir == null)
			dir = new File(".");
		out = new FileOutputStream(dir.getAbsoluteFile() + File.separator + "Main.class");
	}
	
	@Override
	public void inAProgram(AProgram node) {
		cw = new ClassWriter(0);
		cw.visit(V1_4, ACC_PRIVATE + ACC_FINAL, "Main", null, "java/lang/Object", null);
		mv = cw.visitMethod(ACC_PUBLIC + ACC_STATIC + ACC_FINAL, "main", "([Ljava/lang/String;)V", null, null);
		updateLocalsSize(1);
	}
	
	@Override
	public void outAProgram(AProgram node) {
		mv.visitInsn(RETURN);
		mv.visitMaxs(max_stack_size, max_locals_size);
		mv.visitEnd();
		cw.visitEnd();
		WritableByteChannel byte_channel = Channels.newChannel(out);
		try {
			byte_channel.write(ByteBuffer.wrap(cw.toByteArray()));
		} catch (IOException e) {
			e.printStackTrace();
			System.exit(1);
		}
	}

	@Override
	public void inAPutStmt(APutStmt node) {
		mv.visitFieldInsn(GETSTATIC, "java/lang/System", "out", "Ljava/io/PrintStream;");
	}
	
	@Override
	public void outAPutStmt(APutStmt node) {
		switch (last_type) {
			case Integer:
				mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/PrintStream", "println", "(I)V");
				break;
			case Boolean:
				mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/PrintStream", "println", "(Z)V");
				break;
			case String:
				mv.visitMethodInsn(INVOKEVIRTUAL, "java/io/PrintStream", "println", "(Ljava/lang/String;)V");
				break;
			default: // Possible error handling should go here
				break;
		}
		updateStackSize(5); // one for GETSTATIC, others for possible expr trees (which is 4 max)
	}

	@Override
	public void outAInferDeclStmt(AInferDeclStmt node) {
		TIdentifier first = node.getIdentifier();
		Variable var = (Variable) table.getInfo(first.getText().trim());
		allocateVariable(var);
		for (TIdentifier ident : node.getRest())
			allocateVariable((Variable) table.getInfo(ident.getText().trim()));
	}
	
	@Override
	public void outAInferInitStmt(AInferInitStmt node) {
		TIdentifier first = node.getIdentifier();
		Variable var = (Variable) table.getInfo(first.getText().trim());
		allocateVariable(var);
		mv.visitInsn(DUP);
		storePushed(var);  // POPs
		for (TIdentifier ident : node.getRest()) {
			var = (Variable) table.getInfo(ident.getText().trim());
			allocateVariable(var);
			mv.visitInsn(DUP);
			storePushed(var); // POPs
		}
		mv.visitInsn(POP);
		updateStackSize(2); // DUP duplicates, thus size required is 2
	}

	@Override
	public void outATypedDeclStmt(ATypedDeclStmt node) {
		TIdentifier first = node.getIdentifier();
		allocateVariable((Variable) table.getInfo(first.getText().trim()));
		for (TIdentifier ident : node.getRest())
			allocateVariable((Variable) table.getInfo(ident.getText().trim()));
	}
	
	@Override
	public void outATypedInitStmt(ATypedInitStmt node) {
		TIdentifier first = node.getIdentifier();
		Variable var = (Variable) table.getInfo(first.getText().trim());
		allocateVariable(var);
		mv.visitInsn(DUP);
		storePushed(var);  // POPs
		for (TIdentifier ident : node.getRest()) {
			var = (Variable) table.getInfo(ident.getText().trim());
			allocateVariable(var);
			mv.visitInsn(DUP);
			storePushed(var); // POPs
		}
		mv.visitInsn(POP);
		updateStackSize(2); // DUP duplicates, thus size required is 2
	}
	
	@Override
	public void outAExprStmt(AExprStmt node) {
		// expr stmts pushes 1 value onto stack, so POP it
		if (toLong)
			mv.visitInsn(POP2); // longs occupy two consecutive slots
		else
			mv.visitInsn(POP);
	}
	
	@Override
	public void outAAssignExpr(AAssignExpr node) {
		TIdentifier ident = node.getIdentifier();
		Variable var = (Variable) table.getInfo(ident.getText().trim());
		mv.visitInsn(DUP);
		storePushed(var); // POPs
		updateStackSize(2);
		last_type = var.getType();
	}
	
	@Override
	public void caseABinOprExpr(ABinOprExpr node) {
		final MethodVisitor mv = this.mv;
		AnalysisAdapter caseOperator = new AnalysisAdapter() {
			@Override
			public void caseAAddOpr(AAddOpr node) {
				mv.visitInsn(IADD);
			}

			@Override
			public void caseASubOpr(ASubOpr node) {
				mv.visitInsn(ISUB);
			}

			@Override
			public void caseADivOpr(ADivOpr node) {
				mv.visitInsn(IDIV);
			}

			@Override
			public void caseAIdivOpr(AIdivOpr node) {
				mv.visitInsn(IDIV);
			}

			@Override
			public void caseAMulOpr(AMulOpr node) {
				mv.visitInsn(IMUL);
			}

			@Override
			public void caseAShlOpr(AShlOpr node) {
				mv.visitInsn(ISHL);
			}

			@Override
			public void caseAShrOpr(AShrOpr node) {
				mv.visitInsn(ISHR);
			}

			@Override
			public void caseAEqOpr(AEqOpr node) {
				handleBooleanInsn(IF_ICMPEQ);
			}
			
			@Override
			public void caseANoteqOpr(ANoteqOpr node) {
				handleBooleanInsn(IF_ICMPNE);
			}

			@Override
			public void caseAGteqOpr(AGteqOpr node) {
				handleBooleanInsn(IF_ICMPGE);
			}

			@Override
			public void caseAGtOpr(AGtOpr node) {
				handleBooleanInsn(IF_ICMPGT);
			}

			@Override
			public void caseALteqOpr(ALteqOpr node) {
				handleBooleanInsn(IF_ICMPLE);
			}

			@Override
			public void caseALtOpr(ALtOpr node) {
				handleBooleanInsn(IF_ICMPLT);
			}

			@Override
			public void caseAModOpr(AModOpr node) {
				mv.visitInsn(IREM);
			}

			@Override
			public void caseAOrOpr(AOrOpr node) {
				// handled elsewhere
			}

			@Override
			public void caseAAndOpr(AAndOpr node) {
				// handled elsewhere
			}

			@Override
			public void caseAPowerOpr(APowerOpr node) {
				toLong = true;
				mv.visitMethodInsn(INVOKESTATIC, "java/lang/Math", "pow", "(DD)D");
				mv.visitInsn(D2I);
				updateStackSize(4);
			}
		};
		POpr opr = node.getOpr();
		Label lbl_success = null, lbl_next = null;
		
		if (opr instanceof AAndOpr || opr instanceof AOrOpr) {
			lbl_success = new Label();
			lbl_next = new Label();
		}
		
		if (node.getLeft() != null) {
			node.getLeft().apply(BytecodeGenerator.this);
			if (opr instanceof AOrOpr)
				mv.visitJumpInsn(IFNE, lbl_success); // if top value on stack is TRUE
			if (opr instanceof AAndOpr)
				mv.visitJumpInsn(IFEQ, lbl_success); // if top value on stack is FALSE
		}
		
		if (node.getRight() != null) {
			node.getRight().apply(BytecodeGenerator.this);
			if (opr instanceof AOrOpr)
				mv.visitJumpInsn(IFNE, lbl_success); // if top value on stack is TRUE
			if (opr instanceof AAndOpr)
				mv.visitJumpInsn(IFEQ, lbl_success); // if top value on stack is FALSE
		}
		
		opr.apply(caseOperator);
		
		if (opr instanceof AOrOpr) {
			mv.visitInsn(ICONST_0);				 // PUSH FALSE
			mv.visitJumpInsn(GOTO, lbl_next);	 // GOTO Next
			mv.visitLabel(lbl_success);			 // Success...
			mv.visitInsn(ICONST_1);				 // PUSH TRUE
			mv.visitLabel(lbl_next);			 // Next...
		}
		
		if (opr instanceof AAndOpr) {
			mv.visitInsn(ICONST_1);				 // PUSH TRUE
			mv.visitJumpInsn(GOTO, lbl_next);    // GOTO Next
			mv.visitLabel(lbl_success);          // Success...
			mv.visitInsn(ICONST_0);				 // PUSH FALSE
			mv.visitLabel(lbl_next);		     // Next..
		}
		
		OperatorTypeAnalyzer caseOperatorType = new OperatorTypeAnalyzer(); 
		node.getOpr().apply(caseOperatorType);
		last_type = caseOperatorType.result;
		updateStackSize(2);
		toLong = false;
	}
	
	@Override
	public void outAIdentifierExpr(AIdentifierExpr node) {
		Variable var = (Variable) table.getInfo(node.getIdentifier().getText().trim());
		if (toLong) mv.visitInsn(I2D);
		loadVariable(var);
		last_type = var.getType();
	}
	
	@SuppressWarnings("unchecked")
	@Override
	public void outANumberExpr(ANumberExpr node) {
		Literal<Integer> int_lit = (Literal<Integer>) table.getInfo(node.getNumericLit().getText().trim());
		mv.visitLdcInsn(int_lit.getValue());
		if (toLong) mv.visitInsn(I2D);
		updateStackSize(1);
		last_type = int_lit.getType();
	}
	
	@SuppressWarnings("unchecked")
	@Override
	public void outABooleanExpr(ABooleanExpr node) {
		Literal<Boolean> bool = (Literal<Boolean>) table.getInfo(node.getBooleanLit().getText().trim());
		if (bool.getValue())
			mv.visitInsn(ICONST_1);  // TRUE
		else
			mv.visitInsn(ICONST_0);  // FALSE
		last_type = bool.getType();
	}

	@SuppressWarnings("unchecked")
	@Override
	public void outAStringExpr(AStringExpr node) {
		Literal<String> str = (Literal<String>) table.getInfo(node.getStringLit().getText().trim());
		mv.visitLdcInsn(str.getValue());
		updateStackSize(1);
		last_type = str.getType();
	}
	
	// Private Helper methods
	
	private void updateStackSize(int size) {
		if (size > max_stack_size)
			max_stack_size = size;
	}
	
	private void updateLocalsSize(int size) {
		if (size > max_locals_size)
			max_locals_size = size;
	}
	
	private void allocateVariable(Variable var) {
		var.setId(max_locals_size);
		updateLocalsSize(max_locals_size + 1);
	}
	
	// pops 2 and push 1 to stack
	private void handleBooleanInsn(int jump_insn) {
		Label lbl_iftrue = new Label();
		Label lbl_next = new Label();
		
		mv.visitJumpInsn(jump_insn, lbl_iftrue); // jump if true
		mv.visitInsn(ICONST_0); // FALSE
		mv.visitJumpInsn(GOTO, lbl_next);
		
		mv.visitLabel(lbl_iftrue); // establish lbl_iftrue jump point
		mv.visitInsn(ICONST_1);  // TRUE
		
		mv.visitLabel(lbl_next); // establish lbl_next jump point
	}
	
	// pops 1 values from stack
	private void storePushed(Variable var) {
		int id = var.getId();
		if (var.getType().equals(TypedSymbol.Type.Integer))
			mv.visitVarInsn(ISTORE, id);
		else if (var.getType().equals(TypedSymbol.Type.Boolean))
			mv.visitVarInsn(ISTORE, id);
		else if (var.getType().equals(TypedSymbol.Type.String))
			mv.visitVarInsn(ASTORE, id);
		updateLocalsSize(id + 1); // id's start from 0
	}
	
	// pushes 1 value to stack
	private void loadVariable(Variable var) {
		int id = var.getId();
		if (var.getType().equals(TypedSymbol.Type.Integer))
			mv.visitVarInsn(ILOAD, id);
		if (var.getType().equals(TypedSymbol.Type.Boolean))
			mv.visitVarInsn(ILOAD, id);
		if (var.getType().equals(TypedSymbol.Type.String))
			mv.visitVarInsn(ALOAD, id);
		updateStackSize(1);
	}
}
